import cv2
import numpy as np
import matplotlib.pyplot as plt

def reassign_value( y1, y2, x1, x2 , coor ):
    #compute the coordinate of the center point
    center_x = ( x1 + x2 )//2
    center_y = ( y1 + y2 )//2

    # coor = ( y , x )
    y = coor[0]
    x = coor[1]

    set_distance = max( abs( center_x - x ) , abs( center_y - y ) )
    set_distance = abs( ( x1 - x2 )//2 ) - set_distance

    return set_distance

def lsigmoid( value , alpha ):
    return 1./( 1. + np.exp( - alpha * value ) )

def windows_fn( img_size , \
                y1, y2, x1, x2 , \
                alpha , \
                push  = 0 ):
    """
    this function will construct a square mask window 

    the edge of the square will be set to zero point of learnable
    sigmoid function
    """
    height = img_size[0]
    width  = img_size[1]

    mask = np.zeros( ( height , width ) , dtype = "float32" )

    for h in range( height ):
        for w in range( width ):
            reset_value = reassign_value( y1, y2, x1, x2, ( h, w ) )
            reset_value = reset_value - push
            mask[ h, w ] = lsigmoid( reset_value , alpha )

    #plt.imshow( mask )
    #plt.show()

    return mask

def img_warp( LU, RU, RD, LD, img ):
    """
    this funciton accepts 4 auxiliary points in a order:
        LU( left up ), RU( right up ), RD( right down ), LD( left down )
    """
    # find external tangent rectangle
    left = min( LU[0], RU[0], RD[0], LD[0] )
    righ = max( LU[0], RU[0], RD[0], LD[0] )
    up   = min( LU[1], RU[1], RD[1], LD[1] )
    do   = max( LU[1], RU[1], RD[1], LD[1] )

    w = righ - left + 1
    h = do - up + 1

    warp_LU = [ LU[0] - left , LU[1] - up ]
    warp_RU = [ RU[0] - left , RU[1] - up ]
    warp_RD = [ RD[0] - left , RD[1] - up ]
    warp_LD = [ LD[0] - left , LD[1] - up ]

    rows = img.shape[0]
    cols = img.shape[1]

    img_LU = [ 0, 0 ]
    img_RU = [ cols -1 , 0 ]
    img_RD = [ cols -1 , rows -1 ]
    img_LD = [ 0 , rows -1 ]

    src_list = np.float32( [ img_LU, img_RU, img_RD, img_LD ] )
    dst_list = np.float32( [ warp_LU, warp_RU, warp_RD, warp_LD ] )

    M = cv2.getPerspectiveTransform( src_list , dst_list )
    img_trans = cv2.warpPerspective( img, M, ( w , h ) )

    return img_trans

def img_synthesis( img_wrapper , img_content , bbox ):

    LU = bbox[0]
    RU = bbox[1]
    RD = bbox[2]
    LD = bbox[3]

    # find the externally tangent rectangle of the bbox
    
    left = min( LU[0], RU[0], RD[0], LD[0] )
    righ = max( LU[0], RU[0], RD[0], LD[0] )
    up   = min( LU[1], RU[1], RD[1], LD[1] )
    do   = max( LU[1], RU[1], RD[1], LD[1] )

    w = righ - left + 1
    h = do - up + 1

    square_length = max( w , h )
    square = windows_fn( ( square_length , square_length ), \
                0 , square_length , 0 , square_length , \
                1 , \
                push  = 0 )

    img_content_warp = img_warp( LU, RU, RD, LD , img_content )
    mask_warp = img_warp( LU, RU, RD, LD, square )

    mask_warp = np.reshape( mask_warp , \
            ( mask_warp.shape[0] , mask_warp.shape[1] , 1 ) )

    negative_mask_warp = 1. - mask_warp

    img_content_warp = img_content_warp * mask_warp

    img_wrapper[ up : do+1 , left : righ+1 ] = img_wrapper[ up : do+1 , left : righ+1 ] * \
            negative_mask_warp + mask_warp * img_content_warp

    return img_wrapper

def get_random_wrapper( img, bbox , pts ):
    rows = img.shape[0]
    cols = img.shape[1]

    LU = bbox[0]
    RU = bbox[1]
    RD = bbox[2]
    LD = bbox[3]

    left = min( LU[0], RU[0], RD[0], LD[0] )
    righ = max( LU[0], RU[0], RD[0], LD[0] )
    up   = min( LU[1], RU[1], RD[1], LD[1] )
    do   = max( LU[1], RU[1], RD[1], LD[1] )

    left = max( 0 , left )
    up   = max( 0 , up )
    righ = min( cols -1 , righ )
    do   = min( rows -1 , do )

    Nleft = np.random.randint( low = 0 , high = left+1 )
    Nup   = np.random.randint( low = 0 , high = up+1 )
    Nrigh = np.random.randint( low = righ-1 , high = cols -1 )
    Ndo   = np.random.randint( low = do-1 , high = rows - 1 )

    bb1 = []
    bb2 = []
    for pt in pts:
        x = pt[0] - Nleft + 1
        y = pt[1] - Nup +1 

        bb2.append( [x,y] )

    for pt in bbox:
        x = pt[0] - Nleft +1 
        y = pt[1] - Nup +1

        bb1.append( [x,y] )

    return img[ Nup  : Ndo , Nleft : Nrigh ] , bb1, bb2

def test():
    img = cv2.imread( "./0000007482_pos_fromICBC.jpg" )

    rows, cols , _ = img.shape

    pt1 = [ 0 , 0 ]
    pt2 = [ 0 , cols- 1 ]
    pt3 = [ rows -1 , 0 ]
    pt4 = [ rows -1 , cols - 1 ]

    dst1 = [ rows * 0.33 , cols * 0.05 ]
    dst2 = [ rows * 0.25 , cols * 0.90 ]
    dst3 = [ rows * 0.70 , cols * 0.20 ]
    dst4 = [ rows * 0.90 , cols * 0.80 ]

    dst1 = [ rows * 0.2  , 0 ]
    dst2 = [ 0 , cols - 1 ]
    dst3 = [ rows - 1 , 0 ]
    dst4 = [ rows * 0.9 , cols * 0.9 ]

    pts = np.float32( [ pt1 , pt2 , pt3 , pt4 ] )
    dsts= np.float32( [ dst1, dst2, dst3, dst4] )

    M = cv2.getPerspectiveTransform( pts , dsts )

    img_trans = cv2.warpPerspective( img , M , ( 3* cols//2 , 3* rows//2) )

    plt.subplot( 121 )
    plt.imshow( img_trans )

    plt.subplot( 122 )
    plt.imshow( img )
    plt.show()

if __name__ == "__main__":
    #test()
    i = windows_fn( (300, 300) , 0,300,0,300 , 0.2 , push = 10 )

    img = cv2.imread( "./0000007482_pos_fromICBC.jpg" )

    LU = [ 120, 60 ]
    RU = [ 280, 40 ]
    RD = [ 450, 250 ]
    LD = [ 100, 180 ]

    img1 = img_warp( LU, RU, RD, LD, i )

    img2 = img_warp( LU, RU, RD, LD, img )

    img2 = img2/255.

    img1 = np.reshape( img1 , ( img1.shape[0] , img1.shape[1] , 1 ) )

    img3 = img1 * img2 

    img1 = np.reshape( img1 , ( img1.shape[0] , img1.shape[1] ) )

    print( img1.shape )
    plt.imshow( img3 )
    plt.show()
    """

    img1 = img_warp( LU, RU, RD, LD, img )

    print( img1.shape )

    plt.imshow( img1 )
    plt.show()
    """
